{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jYysdyb-CaWM"
   },
   "source": [
    "# Ungraded lab: Serve a model with TensorFlow Serving\n",
    "------------------------\n",
    "\n",
    "\n",
    "In this lab you will be taking a look at TFX's model serving system for production called [Tensorflow Serving](https://www.tensorflow.org/tfx/guide/serving). This system is highly integrated with the Tensorflow stack and provides an easy and straightforward way of deploying models.\n",
    "\n",
    "Specifically you will:\n",
    "- Learn how to install TF serving.\n",
    "- Load a pretrained model that classifies dogs, birds and cats.\n",
    "- Save it following the conventions needed by TF serving.\n",
    "- Spin a web server using TF serving that will accept requests through HTTP.\n",
    "- Interact with your model via a REST API.\n",
    "- Learn about model versioning.\n",
    "\n",
    "This lab draws inspiration from [this](https://www.tensorflow.org/tfx/tutorials/serving/rest_simple) official Tensorflow tutorial, so check it out if you got doubts about the topics covered here.\n",
    "\n",
    "Notice that unlike last ungraded lab, you will be working with TF serving without using Docker. This is to show you different ways in which this serving system can be used.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cXp_5RW_dO7k"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dzLKpmZICaWN"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "import subprocess\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "from IPython.display import Image, display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Guf4H620nSWU"
   },
   "source": [
    "### Downloading the data\n",
    "\n",
    "During this lab you are not going to train a model but to make use of one to get predictions, because of this you need some test images. The model you are going to use was originally trained using images from the datasets `cats and dogs` and `caltech birds`. To ask for predictions some images of the test set are provided:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1x02JYzMgfTB"
   },
   "outputs": [],
   "source": [
    "# Download the images\n",
    "!wget -q https://storage.googleapis.com/mlep-public/course_3/week2/images.zip\n",
    "\n",
    "# Set a base directory\n",
    "base_dir = '/tmp/data'\n",
    "\n",
    "# Unzip images\n",
    "with zipfile.ZipFile('/content/images.zip', 'r') as my_zip:\n",
    "  my_zip.extractall(base_dir)\n",
    "\n",
    "# Save paths for images of each class\n",
    "dogs_dir = os.path.join(base_dir, 'images/dogs')\n",
    "cats_dir = os.path.join(base_dir,'images/cats')\n",
    "birds_dir = os.path.join(base_dir,'images/birds')\n",
    "\n",
    "# Print number of images for each class\n",
    "print(f\"There are {len(os.listdir(dogs_dir))} images of dogs\")\n",
    "print(f\"There are {len(os.listdir(cats_dir))} images of cats\")\n",
    "print(f\"There are {len(os.listdir(birds_dir))} images of birds\\n\\n\")\n",
    "\n",
    "# Look at sample images of each class\n",
    "print(\"Sample cat image:\")\n",
    "display(Image(filename=f\"{os.path.join(cats_dir, os.listdir(cats_dir)[0])}\"))\n",
    "print(\"\\nSample dog image:\")\n",
    "display(Image(filename=f\"{os.path.join(dogs_dir, os.listdir(dogs_dir)[0])}\"))\n",
    "print(\"\\nSample bird image:\")\n",
    "display(Image(filename=f\"{os.path.join(birds_dir, os.listdir(birds_dir)[0])}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0aEMMrP1oyfx"
   },
   "source": [
    "Now that you are familiar with the data you're going to be working with, let's jump to the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5jAk1ZXqTJqN"
   },
   "source": [
    "### Load a pretrained model\n",
    "\n",
    "The purpose of this lab is to showcase TF serving's capabilities so you are not going to spend any time training a model. Instead you will be using a model that you trained during course 1 of the specialization. This model classifies images of birds, cats and dogs and has been trained with image augmentation so it yields really good results.\n",
    "\n",
    "First, download the necessary files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AJuNJRCDdR6J"
   },
   "outputs": [],
   "source": [
    "!wget -q -P /content/model/ https://storage.googleapis.com/mlep-public/course_1/week2/model-augmented/saved_model.pb\n",
    "!wget -q -P /content/model/variables/ https://storage.googleapis.com/mlep-public/course_1/week2/model-augmented/variables/variables.data-00000-of-00001\n",
    "!wget -q -P /content/model/variables/ https://storage.googleapis.com/mlep-public/course_1/week2/model-augmented/variables/variables.index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iv3Rlpvde9g9"
   },
   "source": [
    "Now, load the model into memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vQX6BPAVdiwg"
   },
   "outputs": [],
   "source": [
    "model = tf.keras.models.load_model('/content/model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LAjBnjuqfGuS"
   },
   "source": [
    "At this point you can assume you have succesfully trained the model yourself. You can ignore the warnings about the model being trained on an older version of TensorFlow.\n",
    "\n",
    "For context, ths model uses a simple CNN architecture. Take a quick look at the layers that made it up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yWkoEffIdwLQ"
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AwGPItyphqXT"
   },
   "source": [
    "## Save your model\n",
    "\n",
    "To load our trained model into TensorFlow Serving we first need to save it in [SavedModel](https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/saved_model) format.  This will create a [protobuf](https://developers.google.com/protocol-buffers) file in a well-defined directory hierarchy, and will include a version number.  [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving) allows us to select which version of a model, or \"servable\" we want to use when we make inference requests.  Each version will be exported to a different sub-directory under the given path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0w5Rq8SsgWE6"
   },
   "outputs": [],
   "source": [
    "# Fetch the Keras session and save the model\n",
    "# The signature definition is defined by the input and output tensors,\n",
    "# and stored with the default serving key\n",
    "import tempfile\n",
    "\n",
    "MODEL_DIR = tempfile.gettempdir()\n",
    "version = 1\n",
    "export_path = os.path.join(MODEL_DIR, str(version))\n",
    "print(f'export_path = {export_path}\\n')\n",
    "\n",
    "\n",
    "# Save the model\n",
    "tf.keras.models.save_model(\n",
    "    model,\n",
    "    export_path,\n",
    "    overwrite=True,\n",
    "    include_optimizer=True,\n",
    "    save_format=None,\n",
    "    signatures=None,\n",
    "    options=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wLieWFZzqzoC"
   },
   "source": [
    "A saved model on disk includes the following files:\n",
    "- `assets`: a directory including arbitrary files used by the TF graph.\n",
    "- `variables`: a directory containing information about the training checkpoints of the model.\n",
    "- `saved_model.pb`: the protobuf file that represents the actual TF program.\n",
    "\n",
    "Take a quick look at these files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V-1ARuKsqDpN"
   },
   "outputs": [],
   "source": [
    "print(f'\\nFiles of model saved in {export_path }:\\n')\n",
    "!ls -lh {export_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FM7B_RuDYoIj"
   },
   "source": [
    "## Examine your saved model\n",
    "\n",
    "We'll use the command line utility `saved_model_cli` to look at the [MetaGraphDefs](https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/MetaGraphDef) (the models) and [SignatureDefs](https://www.tensorflow.org/tfx/serving/signature_defs) (the methods you can call) in our SavedModel.  See [this discussion of the SavedModel CLI](https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/saved_model.md#cli-to-inspect-and-execute-savedmodel) in the TensorFlow Guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LU4GDF_aYtfQ"
   },
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {export_path} --tag_set serve --signature_def serving_default"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lSPWuegUb7Eo"
   },
   "source": [
    "That tells us a lot about our model!  In this case we didn't explicitly train the model, so any information about the inputs and outputs is very valuable. For instance we know that this model expects our inputs to be of shape `(150, 150, 3)`, which in combination with the use of `conv2d` layers suggests this model expects colored images in a resolution of `150 by 150`. Also the output of the model are of shape `(3)` suggesting a `softmax` activation with 3 classes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H2IJklPdvb4q"
   },
   "source": [
    "## Prepare data for inference\n",
    "\n",
    "Now that you know the shape of the data expected by the model it is time to preprocess the test images accordingly. These images come in a wide variety of resolutions, luckily Keras has you covered with its [`ImageDataGenerator`](https://keras.io/api/preprocessing/image/). Using this object you can:\n",
    "- Normalize pixel values.\n",
    "- Standardize image resolutions.\n",
    "- Set a batch size for inference.\n",
    "- And more!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UcnU-SjhvYEX"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "\n",
    "# Normalize pixel values\n",
    "test_datagen = ImageDataGenerator(rescale=1./255)\n",
    "\n",
    "# Point to the directory with the test images\n",
    "val_gen_no_shuffle = test_datagen.flow_from_directory(\n",
    "    '/tmp/data/images',\n",
    "    target_size=(150, 150),\n",
    "    batch_size=32,\n",
    "    class_mode='binary',\n",
    "    shuffle=True)\n",
    "\n",
    "# Print the label that is assigned to each class\n",
    "print(f\"labels for each class in the test generator are: {val_gen_no_shuffle.class_indices}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kvw4doL20UeK"
   },
   "source": [
    "Since this object is a generator, you can get a batch of images and labels using the `next` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aGaajm_F0Jd8"
   },
   "outputs": [],
   "source": [
    "# Get a batch of 32 images along with their true label\n",
    "data_imgs, labels = next(val_gen_no_shuffle)\n",
    "\n",
    "# Check shapes\n",
    "print(f\"data_imgs has shape: {data_imgs.shape}\")\n",
    "print(f\"labels has shape: {labels.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3MH3YPjZ1D8I"
   },
   "source": [
    "As expected `data_imgs` is an array containing 32 colored images of 150x150 resolution. In a similar fashion `labels` has the true label for each one of these 32 images.\n",
    "\n",
    "To check that everything is working properly, do a sanity check to plot the first 5 images in the batch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "31yuK5RTvgkT"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.preprocessing.image import array_to_img\n",
    "\n",
    "# Returns string representation of each class\n",
    "def get_class(index):\n",
    "  if index == 0:\n",
    "    return \"bird\"\n",
    "  elif index == 1:\n",
    "    return \"cat\"\n",
    "  elif index == 2:\n",
    "    return \"dog\"\n",
    "  return None\n",
    "\n",
    "\n",
    "# Plots a numpy array representing an image\n",
    "def plot_array(array, label, pred=None):\n",
    "  array = np.squeeze(array)\n",
    "  img = array_to_img(array)\n",
    "  display(img)\n",
    "  if pred is None:\n",
    "    print(f\"Image shows a {get_class(label)}.\\n\")\n",
    "  else:\n",
    "    print(f\"Image shows a {get_class(label)}. Model predicted it was {get_class(pred)}.\\n\")\n",
    "\n",
    "\n",
    "# Plot the first 5 images in the batch\n",
    "for i in range(5):\n",
    "  plot_array(data_imgs[i], labels[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ze1-Z77W4yU6"
   },
   "source": [
    "All images have the same resolution and the true labels are correct.\n",
    "\n",
    "Let's jump to serving the model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DBgsyhytS6KD"
   },
   "source": [
    "## Serve your model with TensorFlow Serving\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W1ZVp_VOU7Wu"
   },
   "source": [
    "### Install TensorFlow Serving\n",
    "\n",
    "You will need to install an older version (2.8.0) because more recent versions are currently incompatible with Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ygwa9AgRloYy"
   },
   "outputs": [],
   "source": [
    "!wget 'http://storage.googleapis.com/tensorflow-serving-apt/pool/tensorflow-model-server-universal-2.8.0/t/tensorflow-model-server-universal/tensorflow-model-server-universal_2.8.0_all.deb'\n",
    "!dpkg -i tensorflow-model-server-universal_2.8.0_all.deb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k5NrYdQeVm52"
   },
   "source": [
    "### Start running TensorFlow Serving\n",
    "\n",
    "This is where we start running TensorFlow Serving and load our model.  After it loads we can start making inference requests using REST.  There are some important parameters:\n",
    "\n",
    "* `rest_api_port`: The port that you'll use for REST requests.\n",
    "* `model_name`: You'll use this in the URL of REST requests.  It can be anything. For this case `animal_classifier` is used.\n",
    "* `model_base_path`: This is the path to the directory where you've saved your model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aUgp3vUdU5GS"
   },
   "outputs": [],
   "source": [
    "# Define an env variable with the path to where the model is saved\n",
    "os.environ[\"MODEL_DIR\"] = MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kJDhHNJVnaLN"
   },
   "outputs": [],
   "source": [
    "# Spin up TF serving server\n",
    "%%bash --bg \n",
    "nohup tensorflow_model_server \\\n",
    "  --rest_api_port=8501 \\\n",
    "  --model_name=animal_classifier \\\n",
    "  --model_base_path=\"${MODEL_DIR}\" >server.log 2>&1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-G3OJcFx6ts2"
   },
   "source": [
    "Take a look at the end of the logs printed out my TF model server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IxbeiOCUUs2z"
   },
   "outputs": [],
   "source": [
    "!tail server.log"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K8Vat6xJ61rW"
   },
   "source": [
    "The server was able to succesfully load and serve the model!\n",
    "\n",
    "Since you are going to interact with the server through HTTP/REST, you should point the requests to `localhost:8501` as it is being printed in the logs above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vwg1JKaGXWAg"
   },
   "source": [
    "## Make a request to your model in TensorFlow Serving\n",
    "\n",
    "At this point you already know how your test data looks like. You are going to make predictions for colored images of 150x150 in batches of 32 images (represented by numpy arrays) at a time.\n",
    "\n",
    "Since REST expects the data to be in JSON format and JSON does not support custom Python data types such as numpy arrays you first need to convert these arrays into nested lists.\n",
    "\n",
    "TF serving expects a field called `instances` which contains the input tensors for the model. To pass in your data to the model you should create a JSON with your data as value for the key `instances`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2dsD7KQG1m-R"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Convert numpy array to list\n",
    "data_imgs_list = data_imgs.tolist()\n",
    "\n",
    "# Create JSON to use in the request\n",
    "data = json.dumps({\"instances\": data_imgs_list})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ReQd4QESIwXN"
   },
   "source": [
    "### Make REST requests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iT3J-lHrhOYQ"
   },
   "source": [
    "We'll send a predict request as a POST request to our server's REST endpoint, and pass it the batch of 32 images. \n",
    "\n",
    "Remember that the endpoint that serves the model is located at `http://localhost:8501`. However this URL still needs some additional parameters to properly handle the request. You should append `v1/models/name-of-your-model:predict` to it so TF serving knows which model to look for and to perform a predict task.\n",
    "\n",
    "You should also pass to the request the data containing the list that represents the 32 images along with a headers dictionary that specifies the type of content that will be passed, which is JSON in this case.\n",
    "\n",
    "After you get a response from the server you can get the predictions out of it by inspecting the `predictions` field of the JSON that the response returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vGvFyuIzW6n6"
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "# Define headers with content-type set to json\n",
    "headers = {\"content-type\": \"application/json\"}\n",
    "\n",
    "# Capture the response by making a request to the appropiate URL with the appropiate parameters\n",
    "json_response = requests.post('http://localhost:8501/v1/models/animal_classifier:predict', data=data, headers=headers)\n",
    "\n",
    "# Parse the predictions out of the response\n",
    "predictions = json.loads(json_response.text)['predictions']\n",
    "\n",
    "# Print shape of predictions\n",
    "print(f\"predictions has shape: {np.asarray(predictions).shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sH2ZComkrTb_"
   },
   "source": [
    "You might think it is weird that the predictions returned 3 values for each image. However, remember that the last layer of the model is a `softmax` function, so it returned a value for each one of the class. To get the actual predictions you need to find the maximum argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k3cTuqkpp-m-"
   },
   "outputs": [],
   "source": [
    "# Compute argmax\n",
    "preds = np.argmax(predictions, axis=1)\n",
    "\n",
    "# Print shape of predictions\n",
    "print(f\"preds has shape: {preds.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2sUyQF1ysKH4"
   },
   "source": [
    "Now you have a predicted class for each one of the test images! Nice!\n",
    "\n",
    "To test how good the model is performing let's plot the first 10 images along with the true and predicted labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tle4OH2lX6oM"
   },
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "  plot_array(data_imgs[i], labels[i], preds[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xM0EceLpss1Y"
   },
   "source": [
    "To do some further testing you can plot more images out of the 32 or even try to generate a new batch from the generator and repeat the steps above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional Challenge\n",
    "\n",
    "Try to recreating the steps above for the next batch of 32 images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code here\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-8NYLaYpuSEP"
   },
   "source": [
    "## Solution\n",
    "\n",
    "If you want some help, the answer can be found in the next cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vQsIJxyttDqQ"
   },
   "outputs": [],
   "source": [
    "# Get a batch of 32 images along with their true label\n",
    "data_imgs, labels = next(val_gen_no_shuffle)\n",
    "\n",
    "# Convert numpy array to list\n",
    "data_imgs_list = data_imgs.tolist()\n",
    "\n",
    "# Create JSON to use in the request\n",
    "data = json.dumps({\"instances\": data_imgs_list})\n",
    "\n",
    "# Capture the response by making a request to the appropiate URL with the appropiate parameters\n",
    "json_response = requests.post('http://localhost:8501/v1/models/animal_classifier:predict', data=data, headers=headers)\n",
    "\n",
    "# Parse the predictions out of the response\n",
    "predictions = json.loads(json_response.text)['predictions']\n",
    "\n",
    "# Compute argmax\n",
    "preds = np.argmax(predictions, axis=1)\n",
    "\n",
    "for i in range(5):\n",
    "  plot_array(data_imgs[i], labels[i], preds[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z5wsVTEssIep"
   },
   "source": [
    "# Conclusion\n",
    "**Congratulations on finishing this ungraded lab!**\n",
    "\n",
    "Now you should have a deeper understanding of TF serving's internals. In the previous ungraded lab you saw how to use TFS alongside with Docker. In this one you saw how TFS and `tensorflow-model-server` worked on their own. You also saw how to save a model and the structure of a saved model.\n",
    "\n",
    "**Keep it up!**"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "-8NYLaYpuSEP"
   ],
   "name": "C4_W1_Lab_2_TF_Serving.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
